import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
from mpl_toolkits.mplot3d import Axes3D

# 假设 agent_positions 是一个 (A, T, 2) 的数组
# A: 车辆数量, T: 时间步数, 2: 车辆的 (x, y) 坐标
A = 100  # 车辆数量
T = 50   # 时间步数
np.random.seed(0)

# 随机生成车辆的位置数据
agent_positions = np.random.rand(A, T, 2) * 100  # 车辆的位置在 0 到 100 的区域内

# 选择多个时间步 (例如 t=10 到 t=30)，合并这些时间步的数据
start_time = 10
end_time = 30
positions_at_multiple_times = agent_positions[:, start_time:end_time, :].reshape(-1, 2)

# 提取所有时间步中车辆的 (x, y) 坐标
x_multiple = positions_at_multiple_times[:, 0]
y_multiple = positions_at_multiple_times[:, 1]

# 估算车辆位置的均值和标准差
mu_x, mu_y = np.mean(x_multiple), np.mean(y_multiple)  # 计算均值
sigma_x, sigma_y = np.std(x_multiple), np.std(y_multiple)  # 计算标准差

# 打印估算结果
print(f"Estimated Mean: mu_x = {mu_x:.2f}, mu_y = {mu_y:.2f}")
print(f"Estimated Standard Deviation: sigma_x = {sigma_x:.2f}, sigma_y = {sigma_y:.2f}")

# 使用估算的参数创建二维高斯分布
x = np.linspace(0, 100, 100)
y = np.linspace(0, 100, 100)
X, Y = np.meshgrid(x, y)

# 创建二维高斯分布
rv = multivariate_normal([mu_x, mu_y], [[sigma_x**2, 0], [0, sigma_y**2]])
Z = rv.pdf(np.dstack((X, Y)))

# 创建三维图
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 绘制高斯分布的三维表面图
ax.plot_surface(X, Y, Z, cmap='viridis', edgecolor='none')

# 绘制实际车辆位置的散点图
ax.scatter(x_multiple, y_multiple, c='r', marker='o', alpha=0.5)

# 设置标题和标签
ax.set_title('3D Gaussian Distribution of Vehicle Positions')
ax.set_xlabel('X Position')
ax.set_ylabel('Y Position')
ax.set_zlabel('Density')

plt.show()
